import * as tf from "@tensorflow/tfjs"
import * as tfvis from "@tensorflow/tfjs-vis"
import {MnistData} from "./data"

window.onload = async () => {
    // 创建MNIST对象
    const data = new MnistData();
    // 加载数据
    await data.load();

    // 获取数据查看数据结构
    const samples = data.nextTestBatch(20);
    console.log(samples)

    // 创建sutface对象用于显示图片
    const surface = tfvis.visor().surface({name: "输入示例"})
    for(let i=0;i<20;i++){
        const imageTensor = tf.tidy(() => {
            return samples.xs.slice([i, 0], [1, 784]).reshape([28,28,1]);
        });

        // 创建Canvas对象
        const canvas = document.createElement("canvas");
        canvas.width = 28;
        canvas.height = 28;
        // 每张图片外边距4px
        canvas.style = "margin: 4px";
        // 可视化图片
        await tf.browser.toPixels(imageTensor, canvas);
        surface.drawArea.appendChild(canvas)
    }


    // 构建卷积神经网络
    const model = tf.sequential();
    // 添加卷积层
    model.add(tf.layers.conv2d({
        inputShape: [28, 28, 1],
        kernelSize: 5,
        filters: 8,
        strides: 1,
        activation: 'relu',
        kernelInitializer: 'varianceScaling'
    }));
    // 添加最大池化层
    model.add(tf.layers.maxPool2d({
        poolSize: [2 ,2],
        strides: [2, 2]
    }));
    // 添加卷积层
    model.add(tf.layers.conv2d({
        kernelSize: 5,
        filters: 16,
        strides: 1,
        activation: 'relu',
        kernelInitializer: 'varianceScaling'
    }));
    // 添加最大池化层
    model.add(tf.layers.maxPool2d({
        poolSize: [2 ,2],
        strides: [2, 2]
    }));
    // 展平
    model.add(tf.layers.flatten());
    // 全连接层
    model.add(tf.layers.dense({
        units: 10,
        activation: 'softmax',
        kernelInitializer: 'varianceScaling'
    }))

    // 配置损失函数和优化器
    model.compile({
        loss: "categoricalCrossentropy",
        optimizer: tf.train.adam(),
        metrics: 'accuracy'
    });

    // 准备训练集和验证集
    const [train_x, train_y]  = tf.tidy(() => {
        const train_data  = data.nextTrainBatch(5000);
        return [
            // 需要将训练数据成卷积第一层的输入形状
            train_data.xs.reshape([5000, 28, 28, 1]),
            train_data.labels,
        ]
    });

    const [val_x, val_y]  = tf.tidy(() => {
        const val_data  = data.nextTestBatch(1000);
        return [
            // 需要将训练数据成卷积第一层的输入形状
            val_data.xs.reshape([1000, 28, 28, 1]),
            val_data.labels,
        ]
    });

    // 训练模型并可视化训练过程
    await model.fit(train_x, train_y, {
        validationData: [val_x, val_y],
        batchSize: 128,
        epochs: 50,
        callbacks: tfvis.show.fitCallbacks(
            {name: '训练过程'},
            ['loss', 'val_loss', 'acc', 'val_acc'],
            {callbacks: ['onEpochEnd']}
        )
    });


    const canvas = document.querySelector('canvas');
    // 绑定鼠标事件：按住左键移动绘制线条(利用矩阵连起来书写数字)
    canvas.addEventListener("mousemove", (e) => {
        if(e.buttons === 1){
            const ctx = canvas.getContext('2d');
            ctx.fillStyle = 'rgb(255,255,255)',
            ctx.fillRect(e.offsetX,e.offsetY,25,25)
        }
    })

    window.clear = () => {
        const ctx = canvas.getContext('2d');
        ctx.fillStyle = 'rgb(0,0,0)',
        ctx.fillRect(0,0,300,300)
    }

    window.predict = () => {
        // 将canvas转换成Tensor，形状是28*28，黑白图片,并归一化
        const input = tf.tidy(() => {
            return tf.image.resizeBilinear(
                tf.browser.fromPixels(canvas),
                [28,28],
                true,
                )
                .slice([0,0,0], [28,28,1])
                .toFloat()
                .div(255)
                .reshape([1, 28,28,1])
            
        });
        // 预测
        const pred = model.predict(input).argMax(1);
        alert( `预测结果为：${pred.dataSync()[0]}`)
    }

};